import json
import re
import os
import sys


def adjust_answer(single_symbol, answer, big_dict):
    if single_symbol+answer in big_dict:
        return ['', single_symbol+answer]
    if answer == "":
        return ['', answer]
    # use re to match / or |
    if '/' in answer or '|' in answer:
        repaired_answers = list()
        candidated_answers = re.split(r"[\/\|]", answer)
        for one_answer in candidated_answers:
            if single_symbol+one_answer in big_dict:
                repaired_answers.append(single_symbol+one_answer)
            else:
                repaired_answers.append(one_answer)
        return ['', "/".join(repaired_answers)]
    if answer in big_dict and single_symbol != "a":
        return ['', answer]
    return ['', answer]


if __name__ == "__main__":

    data = [json.loads(line.strip())
            for line in open("gram_item_raw.json", 'r').readlines()]
    words = json.load(open("all_words.json", 'r'))
    fw = open("gram_item_processed.json", 'w', encoding='utf-8')

    count = 0
    for d in data:
        question = d['desc']
        answer = d['answer']
        prompt = d['prompt']
        print("question: ", question)
        question_tokens = re.split(r'\s+', question)
        repaired_question = []
        for i, token in enumerate(question_tokens):
            if question_tokens[i] == 'm' and (i-1) > -1 and question_tokens[i-1] == "I":
                repaired_question[-1] = repaired_question[-1]+"'m"
            elif question_tokens[i] == 't' and (i-1) > -1 and question_tokens[i-1] == "don":
                repaired_question[-1] = repaired_question[-1]+"\'t"
            elif token == "s" and (i+1) < len(question_tokens) and not question_tokens[i+1].startswith("[["):
                repaired_question[-1] = repaired_question[-1]+"'s"
            else:
                repaired_question.append(token)

        question_tokens = repaired_question
        for i, token in enumerate(question_tokens):
            # print("token: ", token)
            if re.match(r'[a-z]', question_tokens[i]) and len(question_tokens[i]) == 1 and (i+1 < len(question_tokens) and question_tokens[i+1].startswith("[[")):
                answer_index = eval(
                    re.match(r'\[\[([0-9]+)\]\]', question_tokens[i+1]).group(1))
                repair_result = adjust_answer(
                    question_tokens[i], answer[answer_index], words)
                question_tokens[i] = repair_result[0]
                answer[answer_index] = repair_result[1]

            count += 1
        new_dict = {
            "desc": " ".join(question_tokens),
            "answer": answer,
            "prompt": prompt
        }
        print(new_dict)
        fw.write(json.dumps(new_dict)+"\n")
    print("count: ", count)
